import pandas as pd
import datetime
from utils.log import Logger
from utils.common import data_preprocessing
import joblib
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, classification_report
from utils.log import Logger
from sklearn.metrics import roc_curve, auc
from sklearn.metrics import confusion_matrix
import seaborn as sns
# 导入AUC
from sklearn.metrics import roc_auc_score
from xgboost import XGBClassifier
from xgboost import plot_importance

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15


class PowerLoadPredict(object):
    def __init__(self, filename):
        # 配置日志记录
        logfile_name = "predict_" + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        self.logfile = Logger('../', logfile_name).get_logger()
        # 获取数据源
        self.data_source = data_preprocessing(filename)


if __name__ == '__main__':
    df = PowerLoadPredict('../data/test.csv')
    data = df.data_source

    # 为了预测模型，移除数据集中与预测目标相关性较低或不必要的列
    # 这包括员工的个人标识信息、基本工作满意度、性别、部门及股票期权水平等属性
    x = data.drop(columns=['Attrition', 'EmployeeNumber', 'Over18', 'StandardHours', 'PerformanceRating',
                           'EnvironmentSatisfaction', 'Gender', 'Department', 'StockOptionLevel'])

    # 假设 df 是你的原始 DataFrame，包含特征和标签 Department
    categorical_cols = ['BusinessTravel', 'EducationField', 'JobRole', 'MaritalStatus', 'OverTime']
    # # 对分类变量进行 One-Hot 编码
    x = pd.get_dummies(x, columns=categorical_cols)
    #  分离特征和标签
    x = x.drop(columns=['OverTime_No'], axis=1)
    y = data['Attrition']

    model = joblib.load('../model/xgb.pkl')
    y_pred = model.predict(x)
    accuracy = accuracy_score(y, y_pred)
    print(f"Accuracy: {accuracy:.2f}")
    print(f'分类结果报告: {classification_report(y, y_pred)}')
    df.logfile.info(f'分类结果报告: {classification_report(y, y_pred)}')

    # 计算混淆矩阵并可视化
    cm = confusion_matrix(y, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.title('混淆矩阵')
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    # 可选地保存混淆矩阵图到指定路径
    plt.savefig('../data/fig/混淆矩阵.png')

    # 获取模型预测的正类概率
    y_prob = model.predict_proba(x)[:, 1]
    # 计算 ROC 曲线所需指标
    fpr, tpr, thresholds = roc_curve(y, y_prob)
    roc_auc = auc(fpr, tpr)
    # 绘制 ROC 曲线
    plt.figure()
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], 'k--', lw=2)  # 绘制对角线
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC特征曲线 (ROC Curve)')
    plt.legend(loc="lower right")
    plt.savefig('../data/fig/ROC曲线.png')  # 可选保存路径
    # 求AUC
    auc_score = roc_auc_score(y, y_prob)
    print(f'AUC: {auc_score:.2f}')
    df.logfile.info(f'AUC: {auc_score:.2f}')
